%% Larger Transfers Financed with more Progressive Taxes? On the Optimal Design of Taxes and Transfers
% Global solution for "analytical" model
% Version: Mean reverting uninsurable shocks; insurable shocks

cepath = '~/Documents/MATLAB/compecon'; path([cepath '/CEtools;' cepath '/CEdemos'],path);




%%%%%%%%%% Part from H and T %%%%%%%%% 

% % Uninsurable shocks
% Discertizing as normal
% numel(w) = 100;
% alpha_mean = -vw/(2*(1-rhoa^2));
% alpha_var = vw/(1-rhoa^2);
% [nodes_alpha,weights_alpha] = qnwnorm(numel(w),alpha_mean,alpha_var);

%% Preference
 
% Discretization alpha
                            % exp(alpha)
nodes_alpha=log(w);
                                     % density
nodes_eps_n = 10;
eps_mean = -pr.ve/2; eps_var = pr.ve;

[nodes_eps,weights_eps] = qnwnorm(nodes_eps_n,eps_mean,eps_var);



% Stack all combinations of insurable and uninsurable shocks
shocks_comb = gridmake(nodes_alpha,nodes_eps);
weights_comb = gridmake(pai,weights_eps);
weights_comb = weights_comb(:,1).*weights_comb(:,2);
alpha_mat = reshape(shocks_comb(:,1),numel(w),nodes_eps_n);
eps_mat = reshape(shocks_comb(:,2),numel(w),nodes_eps_n);
weights_mat =  reshape(weights_comb,numel(w),nodes_eps_n);

%% Labor Supply, Consumption, Government Budget

tol = 1e-10;
max_iter = 5000;

ymean = zeros(1,1);
c_mat = zeros(numel(w),nodes_eps_n);
H_mat = zeros(numel(w),nodes_eps_n);
labor_disutility = zeros(numel(w),nodes_eps_n);

n_mat = zeros(numel(w),nodes_eps_n);
y_mat = zeros(numel(w),nodes_eps_n);
y_mat_family = zeros(numel(w));

y_mat_ind = zeros(numel(w),nodes_eps_n);

yhat_mat = zeros(numel(w),nodes_eps_n);
c_mat_fam= zeros(numel(w),nodes_eps_n);



    lambda_min = baseline.lambda-0.1;
    lambda_max = baseline.lambda+0.1;
    for iter_fiscal = 1:max_iter
        if iter_fiscal == max_iter
            disp('Maximum number of iterations reached for lambda')
            disp(['Fiscal system indexed ', num2str(i)])
        end 
        
        lambda = 0.5*(lambda_min+lambda_max);
        
        for j = 1:numel(w)
            
            H_min = 0;
            H_max = 100;
            
            
            for iter_alpha = 1:max_iter
                if iter_alpha == max_iter
                    disp('Maximum number of iterations reached for finding H as a function of alpha')
                    disp(['Fiscal system indexed ', num2str(i)])
                    disp(['alpha indexed ', num2str(j)])
                    disp(['err_H', num2str(err_H)])
                    disp(['H_max H_min ', num2str([H_max H_min])])

                end
                    
                H = 0.5*(H_min+H_max);
                nominator = lambda^(1/pr.sig)*(1-pr.tau)^(1/pr.sig)*exp(nodes_alpha(j)*(1-pr.tau)/pr.sig) *H^(-pr.tau/pr.sig);
                denominator=(lambda * exp(nodes_alpha(j)*(1-pr.tau)) *H^(1-pr.tau) )^(1/pr.sig);
               
                err_H = H - nominator * pr.Om/denominator;
                                
                
                
                if abs(err_H)< tol
                    H_mat(j,:) = H;
                    c_mat(j,:) = lambda*exp(nodes_alpha(j)*(1-pr.tau))*H^(1-pr.tau) ;
                    labor_disutility(j,:) =(nominator/denominator)^(1+pr.sig)*pr.Om/(1+pr.sig);  
                    y_mat(j,:) = exp(nodes_alpha(j))*H;
                    yhat_mat(j,:) = lambda*y_mat(j,:).^(1-pr.tau);
                    y_mat_family(j) = exp(nodes_alpha(j))*H;
                    break
                elseif err_H > 0
                    H_max =H;
                else
                    H_min =H;
                end
            end
           
            
            for k = 1:nodes_eps_n
              n_mat(j,k) =(nominator/denominator)*exp(nodes_eps(k)*1/pr.sig) ;
              y_mat_ind(j,k) =(nominator/denominator)*exp(nodes_eps(k)*1/pr.sig)*exp(nodes_alpha(j)+nodes_eps(k))  ;       
            end
        end
           
    Y = sum(sum(weights_mat.*y_mat(:,:)));
    Yhat = sum(sum(weights_mat.*yhat_mat(:,:)));
    ymean = Y;
         
        err_fiscal = G  +0- (Y-Yhat); % + fiscal(i,2)
   
        if abs(err_fiscal) < tol
            break
        elseif err_fiscal < 0
            lambda_min = lambda;
        else
            lambda_max = lambda;
        end
     end
     
        
        
   

% analytically  (1+pr.sig) /pr.sig* veps+pr.va+1/pr.pareto^2
if min(min(min(c_mat))) < 0
    disp('Consumption negative')
end
if min(min(min(n_mat))) < 0
    disp('Labor negative')
end

%% Welfare
util_mat  =log(c_mat) -labor_disutility;


disp(['Optimal lambda is :', num2str(lambda)])
disp(['Optimal Welfare :', num2str(W)])


%%disp(['   \ \ \ \ \      Welfare $ = ', num2str(W(loc)), '$ ','    \ \ \ \ \    $\tau = ',num2str(fiscal(loc,1)), '$ ', '    \ \ \ \ \     $T = ',num2str(fiscal(loc,2)) , '$ ', '    \ \ \ \ \    $\lambda =', num2str(lambda_grid(loc)), '$ ',])



















%% Check equation (7) in Heathcote and T
% with alpha from normal
%  -1.7334121658 HSV plus transfers  %-1.733673403 (tau =0)  % -1.739673362434843  T=0 

% with alpha from pareto grid
% -1.785218928477222 hsv plus tr
% -1.790168122808884 with tau=0  %  -1.790110851628756
%   -1.805184967834348 with Tr=0   % -1.805146335142211

% With more variances
% tau=0. TR=0.31 Welfare -1.726026715906188      -1.725397149741172    
%  HSV still worse:      -1.743059074346725 %   -1.742582186303624    %
%  -1.729926641974160  % -1.729916467659204
% HSV + Transfers:  -1.721645946493513   %-1.726397786184218 


%With G =0.185
% Tr=0.09796 W -1.750842671186861  
% HSV tau = :0.32161 W= -1.754900188833543

%With G= 
%HSV: 0.32161 W = :0.32161
% Affine better -1.776770599558094